import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import os
from datetime import datetime
import polib
import requests
import json
import re
import configparser
import threading
import queue
import pprint

# 默认配置
DEFAULT_CONFIG = {
    'API': {
        'api_key': 'sEDayzod5HFKh7UGyFqEWCSiZ6TfAge6kNwhTTanR8YbgtI8',
        'base_url': 'https://api.tu-zi.com/v1',
        'model': 'glm-4-flash',
        'temperature': '0.7',
        'max_tokens': '1000'
    },
    'Prompt': {
        'user_prompt': '你是一个专业的翻译助手，请将以下文本翻译成{Language}，保持专业性和准确性。',
        'test_text': '这是一个测试文本，用于验证翻译API是否正常工作。'
    }
}

class POTranslator:
    def __init__(self):
        self.window = tk.Tk()
        self.window.title("PO文件翻译工具")
        self.window.geometry("500x600")

        # 创建消息队列用于线程通信
        self.message_queue = queue.Queue()

        # 获取当前工作目录
        self.current_dir = os.getcwd()

        # 加载配置
        self.config = self.load_config()

        # 翻译控制标志
        self.stop_translation = False

        # 创建界面元素
        self.create_widgets()

        # 启动消息处理
        self.process_messages()

    def load_config(self):
        config = configparser.ConfigParser()
        config_file = os.path.join(self.current_dir, 'config.ini')

        # 设置默认配置
        config.read_dict(DEFAULT_CONFIG)

        # 如果配置文件存在，读取它
        if os.path.exists(config_file):
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    config.read_file(f)
            except UnicodeDecodeError:
                # 如果UTF-8读取失败，尝试使用系统默认编码
                with open(config_file, 'r', encoding='gbk') as f:
                    config.read_file(f)
            except Exception as e:
                messagebox.showerror("错误", f"读取配置文件失败: {str(e)}")
                return config

            # 确保所有默认配置项都存在
            for section in DEFAULT_CONFIG:
                if not config.has_section(section):
                    config.add_section(section)
                for key, value in DEFAULT_CONFIG[section].items():
                    if not config.has_option(section, key):
                        config[section][key] = value
            # stream默认False
            if not config.has_option('API', 'stream'):
                config['API']['stream'] = 'False'

        # 保存配置（确保所有配置项都存在）
        try:
            with open(config_file, 'w', encoding='utf-8') as f:
                config.write(f)
        except Exception as e:
            messagebox.showerror("错误", f"保存配置文件失败: {str(e)}")

        return config

    def save_api_config(self):
        try:
            self.save_config()
        except Exception as e:
            messagebox.showerror("错误", f"保存配置失败: {str(e)}")

    def save_config(self):
        try:
            # 验证温度值
            temp = float(self.temperature_var.get())
            if not 0 <= temp <= 2:
                messagebox.showerror("错误", "温度值必须在0到2之间")
                return

            # 验证最大长度值
            max_tokens = int(self.max_tokens_var.get())
            if not 100 <= max_tokens <= 4000:
                messagebox.showerror("错误", "最大长度必须在100到4000之间")
                return

            # 更新配置对象
            if not self.config.has_section('API'):
                self.config.add_section('API')
            if not self.config.has_section('Prompt'):
                self.config.add_section('Prompt')

            self.config['API']['api_key'] = self.api_key_var.get()
            self.config['API']['base_url'] = self.base_url_var.get()
            self.config['API']['model'] = self.model_var.get()
            self.config['API']['temperature'] = str(temp)
            self.config['API']['max_tokens'] = str(max_tokens)
            self.config['API']['stream'] = str(self.stream_var.get())

            # 获取文本框中的提示词和测试文本
            self.config['Prompt']['user_prompt'] = self.prompt_text.get('1.0', tk.END).strip()
            self.config['Prompt']['test_text'] = self.test_text_var.get()

            # 保存到文件
            config_file = os.path.join(self.current_dir, 'config.ini')
            try:
                with open(config_file, 'w', encoding='utf-8') as f:
                    self.config.write(f)
                messagebox.showinfo("成功", "配置已保存")
            except Exception as e:
                messagebox.showerror("错误", f"保存配置文件失败: {str(e)}")

        except ValueError:
            messagebox.showerror("错误", "请输入有效的数值")
        except Exception as e:
            messagebox.showerror("错误", f"保存配置失败: {str(e)}")

    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.window, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 创建配置框架
        config_frame = ttk.LabelFrame(main_frame, text="API配置", padding="5")
        config_frame.pack(fill=tk.X, pady=5)

        # 打印当前config中的所有信息，便于调试
        for section in self.config.sections():
            print(f"[配置节] {section}")
            for key, value in self.config[section].items():
                print(f"    {key} = {value}")
                
        # API Key
        ttk.Label(config_frame, text="API Key:").grid(row=0, column=0, sticky=tk.W, pady=2)
        self.api_key_var = tk.StringVar()
        self.api_key_var.set(self.config.get('API', 'api_key', fallback=''))
        ttk.Entry(config_frame, textvariable=self.api_key_var, width=50).grid(row=0, column=1, pady=2)

        # Base URL
        ttk.Label(config_frame, text="Base URL:").grid(row=1, column=0, sticky=tk.W, pady=2)
        self.base_url_var = tk.StringVar()
        self.base_url_var.set(self.config.get('API', 'base_url', fallback=''))
        ttk.Entry(config_frame, textvariable=self.base_url_var, width=50).grid(row=1, column=1, pady=2)

        # Model
        ttk.Label(config_frame, text="Model:").grid(row=2, column=0, sticky=tk.W, pady=2)
        self.model_var = tk.StringVar()
        self.model_var.set(self.config.get('API', 'model', fallback=''))
        ttk.Entry(config_frame, textvariable=self.model_var, width=50).grid(row=2, column=1, pady=2)

        # Temperature
        ttk.Label(config_frame, text="Temperature:").grid(row=3, column=0, sticky=tk.W, pady=2)
        temp_frame = ttk.Frame(config_frame)
        temp_frame.grid(row=3, column=1, sticky=tk.W, pady=2)
        self.temperature_var = tk.StringVar()
        print(self.config.get('API', 'temperature', fallback='0.221'))
        self.temperature_var.set(self.config.get('API', 'temperature', fallback='0.221'))
        ttk.Entry(temp_frame, textvariable=self.temperature_var, width=8).pack(side=tk.LEFT, padx=5)
        ttk.Label(temp_frame, text="(0-2)").pack(side=tk.LEFT)

        # Max Tokens
        ttk.Label(config_frame, text="Max Tokens:").grid(row=4, column=0, sticky=tk.W, pady=2)
        tokens_frame = ttk.Frame(config_frame)
        tokens_frame.grid(row=4, column=1, sticky=tk.W, pady=2)
        self.max_tokens_var = tk.StringVar()
        self.max_tokens_var.set(self.config.get('API', 'max_tokens', fallback='1000'))
        ttk.Entry(tokens_frame, textvariable=self.max_tokens_var, width=8).pack(side=tk.LEFT, padx=5)
        ttk.Label(tokens_frame, text="(100-4000)").pack(side=tk.LEFT)

        # stream模式选择
        self.stream_var = tk.BooleanVar()
        self.stream_var.set(self.config.getboolean('API', 'stream', fallback=False))
        stream_check = ttk.Checkbutton(config_frame, text="Ollama流式(stream)模式", variable=self.stream_var)
        stream_check.grid(row=5, column=1, sticky=tk.W, pady=2)

        # 用户提示词
        ttk.Label(config_frame, text="用户提示词:").grid(row=6, column=0, sticky=tk.W, pady=2)

        # 创建多行文本框
        prompt_frame = ttk.Frame(config_frame)
        prompt_frame.grid(row=6, column=1, sticky=tk.W+tk.E, pady=2)

        # 创建文本框和滚动条
        self.prompt_text = tk.Text(prompt_frame, height=4, width=50, wrap=tk.WORD)
        self.prompt_text.delete('1.0', tk.END)
        self.prompt_text.insert('1.0', self.config.get('Prompt', 'user_prompt', fallback=''))
        self.prompt_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 添加滚动条
        scrollbar = ttk.Scrollbar(prompt_frame, orient=tk.VERTICAL, command=self.prompt_text.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.prompt_text['yscrollcommand'] = scrollbar.set

        ttk.Label(config_frame, text="(使用 {Language} 作为语言变量)").grid(row=7, column=1, sticky=tk.W, pady=0)

        # 测试文本输入
        ttk.Label(config_frame, text="测试文本:").grid(row=8, column=0, sticky=tk.W, pady=2)
        self.test_text_var = tk.StringVar()
        self.test_text_var.set(self.config.get('Prompt', 'test_text', fallback=''))
        test_text_entry = ttk.Entry(config_frame, textvariable=self.test_text_var, width=50)
        test_text_entry.grid(row=8, column=1, pady=2)

        # 按钮框架
        button_frame = ttk.Frame(config_frame)
        button_frame.grid(row=9, column=0, columnspan=2, pady=5)

        # 保存配置按钮
        ttk.Button(button_frame, text="保存配置", command=self.save_api_config).pack(side=tk.LEFT, padx=5)

        # 测试翻译按钮
        ttk.Button(button_frame, text="测试翻译", command=self.test_translation).pack(side=tk.LEFT, padx=5)

        # 设置自动保存变量（不显示在GUI上）
        self.auto_save_var = tk.BooleanVar()
        self.auto_save_var.set(True)  # 始终启用自动保存
        self.save_interval = 5  # 固定为每5个条目自动保存一次

        # 文件操作框架
        file_frame = ttk.LabelFrame(main_frame, text="文件操作", padding="5")
        file_frame.pack(fill=tk.X, pady=5)

        # 文件选择按钮
        self.select_btn = ttk.Button(file_frame, text="选择PO文件", command=self.select_file)
        self.select_btn.pack(pady=5)

        # 显示选中文件路径
        self.file_label = ttk.Label(file_frame, text="未选择文件")
        self.file_label.pack(pady=5)

        # 显示目标语言
        self.language_label = ttk.Label(file_frame, text="目标语言: 未选择文件")
        self.language_label.pack(pady=5)

        # 翻译按钮框架
        translate_btn_frame = ttk.Frame(file_frame)
        translate_btn_frame.pack(pady=5)

        # 翻译按钮
        self.translate_btn = ttk.Button(translate_btn_frame, text="开始翻译", command=self.start_translation)
        self.translate_btn.pack(side=tk.LEFT, padx=5)

        # 停止翻译按钮
        self.stop_btn = ttk.Button(translate_btn_frame, text="停止翻译", command=self.stop_translation_process, state=tk.DISABLED)
        self.stop_btn.pack(side=tk.LEFT, padx=5)

        # 进度显示
        self.progress_label = ttk.Label(file_frame, text="")
        self.progress_label.pack(pady=5)

    def select_file(self):
        file_path = filedialog.askopenfilename(filetypes=[("PO files", "*.po")])
        if file_path:
            self.file_label.config(text=file_path)
            self.current_file = file_path
            # 读取并显示目标语言
            target_lang = self.get_target_language(file_path)
            self.language_label.config(text=f"目标语言: {target_lang}")

            # 统计翻译情况
            self.count_translations(file_path)

    def get_target_language(self, file_path):
        try:
            po = polib.pofile(file_path)
            # 从文件头部信息中获取语言代码
            lang_code = None
            for entry in po.metadata:
                if entry == 'Language':
                    lang_code = po.metadata[entry]
                    break

            if lang_code:
                return lang_code
            return "未知语言"
        except Exception as e:
            return f"读取语言信息失败: {str(e)}"

    def count_translations(self, file_path):
        try:
            po = polib.pofile(file_path)
            total = len(po)
            translated = sum(1 for entry in po if entry.msgstr)
            untranslated = total - translated

            # 更新进度显示
            self.progress_label.config(
                text=f"总条目: {total} | 已翻译: {translated} | 未翻译: {untranslated}"
            )
        except Exception as e:
            self.progress_label.config(text=f"统计失败: {str(e)}")

    def build_prompt(self, text, target_lang):
        """统一构建prompt"""
        prompt_template = self.prompt_text.get('1.0', tk.END).strip()
        if '{Info}' in prompt_template:
            return prompt_template.format(Language=target_lang, Info=text)
        else:
            return prompt_template.format(Language=target_lang) + text

    def call_translate_api(self, prompt, is_ollama=False):
        """统一API请求，并打印请求和响应内容"""
        if is_ollama:
            url = f"{self.config['API']['base_url'].rstrip('/')}/api/chat"
            headers = {"Content-Type": "application/json"}
            stream_flag = self.config.getboolean('API', 'stream', fallback=False)
            data = {
                "model": self.config['API']['model'],
                "messages": [
                    {"role": "user", "content": prompt}
                ],
                "stream": stream_flag,
                "options": {
                    "temperature": float(self.config['API']['temperature']),
                    "num_predict": int(self.config['API']['max_tokens'])
                }
            }
            try:
                print("[发送请求] URL:", url)
                print("[发送请求] Headers:")
                pprint.pprint(headers)
                print("[发送请求] Data:")
                pprint.pprint(data)
                response = requests.post(url, headers=headers, json=data, stream=stream_flag)
                print("[收到响应] 状态码:", response.status_code)
                print("[收到响应] 内容:")
                if stream_flag:
                    # 逐行处理流式响应
                    result_text = ""
                    for line in response.iter_lines(decode_unicode=True):
                        if line.strip():
                            print(line)
                            obj = json.loads(line)
                            if "message" in obj and "content" in obj["message"]:
                                result_text += obj["message"]["content"]
                    return result_text
                else:
                    print(response.text)
                    response.raise_for_status()
                    result = response.json()
                    return result['message']['content']
            except Exception as e:
                return f"[Ollama API调用失败] {str(e)}"
        else:
            headers = {
                "Authorization": f"Bearer {self.config['API']['api_key']}",
                "Content-Type": "application/json"
            }
            data = {
                "model": self.config['API']['model'],
                "messages": [
                    {"role": "user", "content": prompt}
                ],
                "temperature": float(self.config['API']['temperature']),
                "max_tokens": int(self.config['API']['max_tokens'])
            }
            try:
                print("[发送请求] URL:", f"{self.config['API']['base_url']}/chat/completions")
                print("[发送请求] Headers:")
                pprint.pprint(headers)
                print("[发送请求] Data:")
                pprint.pprint(data)
                response = requests.post(
                    f"{self.config['API']['base_url']}/chat/completions",
                    headers=headers,
                    json=data
                )
                print("[收到响应] 状态码:", response.status_code)
                print("[收到响应] 内容:")
                print(response.text)
                response.raise_for_status()
                result = response.json()
                return result['choices'][0]['message']['content']
            except Exception as e:
                return f"[翻译API调用失败] {str(e)}"

    def translate_text(self, text, target_lang):
        prompt = self.build_prompt(text, target_lang)
        is_ollama = "127.0.0.1:11434" in self.config['API']['base_url']
        result = self.call_translate_api(prompt, is_ollama)
        if result.startswith('[Ollama API调用失败]') or result.startswith('[翻译API调用失败]'):
            messagebox.showerror("错误", result)
            return None
        return result

    def process_messages(self):
        """处理来自翻译线程的消息"""
        try:
            while True:
                message = self.message_queue.get_nowait()
                if message['type'] == 'progress':
                    self.progress_label.config(text=message['text'])
                elif message['type'] == 'complete':
                    self.progress_label.config(text="翻译完成！")
                    messagebox.showinfo("完成", "翻译已完成并保存")
                    self.translate_btn.config(state=tk.NORMAL)
                    self.stop_btn.config(state=tk.DISABLED)
                elif message['type'] == 'stop_complete':
                    # 翻译已停止，恢复按钮状态
                    self.translate_btn.config(state=tk.NORMAL)
                    self.stop_btn.config(state=tk.DISABLED)
                    messagebox.showinfo("已停止", "翻译已停止，已翻译的内容已保存")
                elif message['type'] == 'error':
                    self.progress_label.config(text="翻译失败")
                    messagebox.showerror("错误", message['text'])
                    self.translate_btn.config(state=tk.NORMAL)
                    self.stop_btn.config(state=tk.DISABLED)
        except queue.Empty:
            pass
        finally:
            self.window.after(100, self.process_messages)

    def start_translation(self):
        # 每次点击前都重新加载配置，确保用最新配置
        self.config = self.load_config()
        if not hasattr(self, 'current_file'):
            messagebox.showwarning("警告", "请先选择PO文件")
            return

        # 获取目标语言
        target_lang = self.get_target_language(self.current_file)
        if target_lang == "未知语言":
            messagebox.showwarning("警告", "无法确定目标语言")
            return

        # 重置停止标志
        self.stop_translation = False

        # 禁用翻译按钮，启用停止按钮
        self.translate_btn.config(state=tk.DISABLED)
        self.stop_btn.config(state=tk.NORMAL)

        # 启动翻译线程
        thread = threading.Thread(target=self.translate_file, args=(self.current_file, target_lang))
        thread.daemon = True
        thread.start()

    def stop_translation_process(self):
        """停止翻译过程"""
        if messagebox.askyesno("确认", "确定要停止翻译吗？已翻译的内容将被保存。"):
            self.stop_translation = True
            self.stop_btn.config(state=tk.DISABLED)
            self.progress_label.config(text="正在停止翻译，请等待当前条目完成...")
            self.window.update()

    def translate_file(self, file_path, target_lang):
        try:
            # 加载PO文件
            po = polib.pofile(file_path)

            # 统计需要翻译的条目
            untranslated = [entry for entry in po if not entry.msgstr]
            total = len(untranslated)

            if total == 0:
                self.message_queue.put({
                    'type': 'complete',
                    'text': '没有需要翻译的条目'
                })
                return

            # 更新进度显示
            self.message_queue.put({
                'type': 'progress',
                'text': f"开始翻译，共 {total} 条待翻译"
            })

            # 使用固定的自动保存设置
            auto_save = self.auto_save_var.get()  # 始终为True
            save_interval = self.save_interval  # 固定为5

            # 翻译每个条目
            last_save_time = datetime.now()
            for i, entry in enumerate(untranslated, 1):
                # 检查是否需要停止翻译
                if self.stop_translation:
                    # 保存当前进度
                    po.save(file_path)
                    self.message_queue.put({
                        'type': 'progress',
                        'text': f"翻译已停止，已完成: {i-1}/{total}，进度已保存"
                    })
                    # 恢复按钮状态
                    self.message_queue.put({
                        'type': 'stop_complete'
                    })
                    return

                # 检查原文是否包含中文字符
                if not self.contains_chinese(entry.msgid):
                    # 如果不包含中文，直接使用原文作为翻译结果
                    entry.msgstr = entry.msgid
                else:
                    # 包含中文，调用翻译API
                    translated = self.translate_text(entry.msgid, target_lang)
                    if translated:
                        entry.msgstr = translated

                # 更新进度
                self.message_queue.put({
                    'type': 'progress',
                    'text': f"正在翻译: {i}/{total}"
                })

                # 自动保存
                if auto_save and i % save_interval == 0:
                    try:
                        # 保存当前进度
                        po.save(file_path)
                        current_time = datetime.now()
                        self.message_queue.put({
                            'type': 'progress',
                            'text': f"正在翻译: {i}/{total} (已自动保存，耗时: {(current_time - last_save_time).total_seconds():.1f}秒)"
                        })
                        last_save_time = current_time
                    except Exception as save_error:
                        self.message_queue.put({
                            'type': 'progress',
                            'text': f"正在翻译: {i}/{total} (自动保存失败: {str(save_error)})"
                        })

            # 保存翻译后的文件
            po.save(file_path)

            # 发送完成消息
            self.message_queue.put({
                'type': 'complete'
            })

        except Exception as e:
            # 发送错误消息
            self.message_queue.put({
                'type': 'error',
                'text': str(e)
            })

    def contains_chinese(self, text):
        """检查文本是否包含中文字符"""
        if not text:
            return False
        chinese_pattern = re.compile('[\u4e00-\u9fff]')
        return bool(chinese_pattern.search(text))

    def test_translation(self):
        test_text = self.test_text_var.get()
        if not test_text.strip():
            messagebox.showwarning("警告", "请输入测试文本")
            return
        prompt = self.build_prompt(test_text, "en")
        is_ollama = "127.0.0.1:11434" in self.config['API']['base_url']
        self.progress_label.config(text="正在测试翻译...")
        self.window.update()
        result = self.call_translate_api(prompt, is_ollama)
        if result.startswith('[Ollama API调用失败]'):
            messagebox.showerror("错误", result)
            self.progress_label.config(text="测试失败")
        elif result.startswith('[翻译API调用失败]'):
            messagebox.showerror("错误", result)
            self.progress_label.config(text="测试失败")
        else:
            messagebox.showinfo("测试结果",
                f"原文：{test_text}\n\n译文：{result}\n\nAPI配置正常！")
            self.progress_label.config(text="测试完成")

    def run(self):
        self.window.mainloop()

if __name__ == "__main__":
    app = POTranslator()
    app.run()